
from vBaum import *
from math import *
from random import *

voxelsPerUnit = 30

rand = Random()
rand.seed(12345)

def Vec3(x, y, z):
	return Vector3(x, z, y)
	
def randVec(sigma):
	x = gauss(0, sigma)
	y = gauss(0, sigma)
	z = gauss(0, sigma)
	return Vec3(x,y,z)
	
def printVec(vec):
	print(str(vec.x) + " " + str(vec.y) + " " + str(vec.z))

# rounds the subvolume of a volume that intersects with the given sphere
def roundSubVolume(volume, sphereCenter, sphereRadius, roundingRadius):
	sphere = Sphere(sphereRadius * voxelsPerUnit)
	sphere.setPosition(sphereCenter * voxelsPerUnit)
	intersection = volume.intersect(sphere)
#	intersection.exportAsObj("treeSphere", voxelsPerUnit, 0)
	volume = volume.subtract(sphere)
	rounded = intersection.roundConcave(roundingRadius * voxelsPerUnit)
#	rounded.exportAsObj("roundedTreeSphere", voxelsPerUnit, 0)
	volume = volume.union(rounded)

	return volume
	
def tesselateLine(p1, p2, numMidsMin, numMidsMax):
	numMids = rand.randint(numMidsMin, numMidsMax)
	offset = 1.0 / (1.0+numMids)
	curr = offset
	mids = []	
	for i in range(0, numMids):
		p = randVec(0.2)	
		ref = (p1*curr) + (p2*(1-curr))
		mids = mids + [ref+p]
		curr = curr + offset
		
	mids = sorted(mids, key = lambda p: p.squaredDistance(p1)) 
	return [p1] + mids + [p2]
	
def scalePoints(points):
	scaledPoints = []
	for p in points:
		scaledPoints.append(p * voxelsPerUnit)
	return scaledPoints

def createBranch(points, radiusBegin, radiusEnd):
	scaledPoints = scalePoints(points)
	scaledRadiusBegin = radiusBegin * voxelsPerUnit
	scaledRadiusEnd = radiusEnd * voxelsPerUnit
	return Spline(scaledPoints, False, lambda x: Sphere(scaledRadiusBegin*(1-x) + scaledRadiusEnd*x))
	
def computeOrthogonalVector(vec):
	someVec = Vector3(1, 1, 1)
	if (abs(vec.x) >= abs(vec.y) and abs(vec.x) >= abs(vec.z)): someVec.x = 0
	elif (abs(vec.y) >= abs(vec.x) and abs(vec.y) >= abs(vec.z)): someVec.y = 0
	else: someVec.z = 0
	someVec.normalise()
	return someVec.crossProduct(vec)
	
def rotateDir(dir, rotateAxis, angle):
	q = Quaternion(angle, rotateAxis)
	return q * dir
	
def computeFirstBranchDir(rootDir, angle):
	ortho = computeOrthogonalVector(rootDir)
	return rotateDir(rootDir, ortho, angle)

# Monopodial: Main axis dominates
class MonopodialBranch:
	def __init__(self):
		self.maxRecursionDepth = 3
		self.branchAngle = 0.8
		self.branchAngleSigma = 0.2
		self.numBranchesPerFork = 3
		self.branchSectionLength = 3
		self.angleBetweenBranchesAtForkSigma = 0.1
		self.divergence = True
		self.successor = self
		self.orientUp = False
		self.RadiusHeightConstant = 25
		self.roundNodes = True
		
	def generateGeometry(self, recursionDepth, startPos, dir, trunkRadius):
		if recursionDepth > self.maxRecursionDepth: return Volume()
		
		maxLength = pow(trunkRadius * trunkRadius, 1.0/3.0) * self.RadiusHeightConstant
		# print("Max length is ", maxLength)
		numForks = floor(maxLength / self.branchSectionLength)
			
		endPos = startPos + dir * maxLength
		points = tesselateLine(startPos, endPos, numForks, numForks)
		branch = createBranch(points, max(trunkRadius, 1.0 / voxelsPerUnit + 0.01), 1.0 / voxelsPerUnit + 0.01)
		
		i = 0
		angleOffset = 0
		for i in range (1, len(points)-1):
			forkPos = points[i]
			forkRadius = (1 - i / (len(points)-1)) * trunkRadius
			angleStep = (2 * pi) / (self.numBranchesPerFork)
			angle = 0
			if self.divergence: angle = angleOffset
			
			for a in range(0, self.numBranchesPerFork):
				firstBranchDir = computeFirstBranchDir(dir, gauss(self.branchAngle, self.branchAngleSigma))
				if self.orientUp:
					firstBranchDir = rotateDir(dir, Vec3(0, 1, 0), gauss(self.branchAngle, self.branchAngleSigma))
				angle = angle + angleStep
				if angle > 2 * pi: angle = angle - (2 * pi)
				branchDir = rotateDir(firstBranchDir, dir, gauss(angle, self.angleBetweenBranchesAtForkSigma))
				subBranch = self.successor.generateGeometry(recursionDepth + 1, forkPos, branchDir, forkRadius * 0.5)
				branch = branch.union(subBranch)
			angleOffset = angleOffset + angleStep * 0.5
			if angleOffset > 2 * pi: angleOffset = angleOffset - (2 * pi)
			
			if self.roundNodes: branch = roundSubVolume(branch, forkPos, forkRadius*2.0, forkRadius*0.5)	
			
		return branch
			
node = MonopodialBranch()
node2 = MonopodialBranch()
node2.branchSectionLength = 1.5
node2.numBranchesPerFork = 2
node2.orientUp = True
node2.orientNodes = False
node2.divergence = False
node.successor = node2
tree = node.generateGeometry(0, Vec3(0, 0, 0), Vec3(0, 1, 0), 0.5)
tree.exportAsObj("monopodialTree", voxelsPerUnit, 5)

# Sympodial: branches dominate
class SympodialBranch:
	def __init__(self, maxRecursionDepth, numSubBranchesMin, numSubBranchesMax, dominantBranchAngle, branchAngle):
		self.maxRecursionDepth = maxRecursionDepth
		self.numSubBranchesMin = numSubBranchesMin
		self.numSubBranchesMax = numSubBranchesMax
		self.dominantBranchAngle = dominantBranchAngle
		self.branchAngle = branchAngle
		self.dominantSubBranchAreaPercent = 0.75
		self.radiusFalloff = 0.75
		
	def generateGeometry(self, recursionDepth, radius, startPos, dir):
		endPos = startPos + dir * gauss(4, 0.5)
		points = tesselateLine(startPos, endPos)
		endRadius = radius * self.radiusFalloff
		branch = createBranch(points, radius, endRadius)
		if (recursionDepth > self.maxRecursionDepth): return Volume()
		
		# fork
		dominantDir = computeFirstBranchDir(dir, gauss(self.dominantBranchAngle, 0.2))
		dominantDirAngle = rand.uniform(0, pi)
		dominantDir = rotateDir(dominantDir, dir, dominantDirAngle)
		totalArea = endRadius * endRadius
		dominantSubBranchArea = totalArea * self.dominantSubBranchAreaPercent
		dominantSubBranch = self.generateGeometry(recursionDepth + 1, sqrt(dominantSubBranchArea), endPos, dominantDir)
		branch = branch.union(dominantSubBranch)
		
		numSubBranches = randint(self.numSubBranchesMin, self.numSubBranchesMax)
		recessiveBranchesArea = totalArea * (1 - self.dominantSubBranchAreaPercent)
		recessiveBranchesRadius = sqrt(recessiveBranchesArea)
		angleStep = (2 * pi) / (numSubBranches + 1)
		angle = dominantDirAngle
		recessiveDir = computeFirstBranchDir(dir, gauss(self.branchAngle, 0.3))
		for a in range(0, numSubBranches):
			angle = angle + angleStep
			if angle > 2 * pi: angle = angle - (2 * pi)
			branchDir = rotateDir(recessiveDir, dir, gauss(angle, 0.1))
			recessiveSubBranch = self.generateGeometry(recursionDepth + 1, recessiveBranchesRadius, endPos, branchDir)
			branch = branch.union(recessiveSubBranch)
			
		branch = roundSubVolume(branch, endPos, endRadius*2.0, endRadius*0.5)
			
		return branch

'''trunkPoints = tesselateLine(Vec3(0, 0, 0), Vec3(0, 5, 0))
branchPoints1 = tesselateLine(Vec3(0, 5, 0), Vec3(4, 8, 3))
branchPoints2 = tesselateLine(Vec3(0, 5, 0), Vec3(-1, 9, -2))

trunk = createBranch(trunkPoints, 1.0, 0.5)
branch1 = createBranch(branchPoints1, 0.5, 0.25)
branch2 = createBranch(branchPoints2, 0.5, 0.25)

tree = trunk.union(branch1)
tree = tree.union(branch2)
tree = roundSubVolume(tree, Vec3(0, 5, 0), 1.0)
tree.exportAsObj("tree", voxelsPerUnit, 0)'''
